/******************************************************************************\
    SocketStream.hpp contains a class of buffered socket.
    Copyright (C) 2008  Victor Vinogradov (fly3333@gmail.com)

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
\******************************************************************************/

#ifndef __SOCKET_STREAM_HPP
#define __SOCKET_STREAM_HPP

#include <winsock2.h>
#include <string>
#include ".\Memory.hpp"
#include ".\KMP.h"
#include ".\CriticalSection.hpp"

class cWardSocketStream
{
	SOCKET Socket;
	cMemory< char > Buffer;
	int BufferOffset, BufferSize;
	size_t MaxBufferSize;
	cCriticalSection ReceiveCriticalSection;
	cCriticalSection SendCriticalSection;
public:
	cWardSocketStream( SOCKET socket, size_t max_buffer_size )
		: Socket( socket )
		, Buffer( max_buffer_size )
		, BufferOffset( 0 )
		, BufferSize( 0 )
		, MaxBufferSize( max_buffer_size )
	{
	}
	int Get()
	{
		cLocker locker( ReceiveCriticalSection );
		if( BufferOffset >= BufferSize )
		{
			BufferOffset = 0;
			BufferSize = recv( Socket, Buffer, ( int )MaxBufferSize, 0 );
			if( BufferSize == SOCKET_ERROR ||
				 BufferSize == 0 )
			{
				BufferSize = 0;
				return -1;
			}
		}
		return ( unsigned char )Buffer[ BufferOffset++ ];
	}
	std::string Get( int size )
	{
		cLocker locker( ReceiveCriticalSection );
		std::string result;
		for( int i = 0; i < size; ++i )
		{
			int c = Get();
			if( c == -1 )
				break;
			result += ( char )c;
		}
		return result;
	}
	bool Eof() const
	{
		return BufferSize == 0;
	}
	std::string GetLine( const std::string &delim = "\r\n", int max_length = -1 )
	{
		cLocker locker( ReceiveCriticalSection );
		std::string result;
		cMemory< int > next( delim.length() + 1 );
		PreKMP( delim.c_str(), ( int )delim.length(), next );
		for( int delim_i = 0, ch, length = 0; delim_i < ( int )delim.length() && length != max_length && ( ch = Get() ) != -1; ++delim_i, ++length )
		{
			result.push_back( ( char )ch );
			while( delim_i >= 0 && ( char )ch != delim[ delim_i ] )
				delim_i = next[ delim_i ];
		}
		return result;
	}
	bool PutLine( const std::string &str )
	{
		cLocker locker( SendCriticalSection );
		for( int sent_length = 0, curr_length; sent_length < ( int )str.length(); sent_length += curr_length )
		{
			curr_length = send( Socket, str.data() + sent_length, ( int )str.length() - sent_length, 0 );
			if( curr_length == SOCKET_ERROR )
				return false;
		}
		return true;
	}
	SOCKET GetSocket() const 
	{
		return Socket;
	}
};

class cSocketStream : public cPointerObject< cWardSocketStream >
{
public:
	bool Create( SOCKET socket, unsigned buffer_size )
	{
		cWardSocketStream *ward_socket_stream;
		try
		{
			ward_socket_stream = new cWardSocketStream( socket, buffer_size );
		}
		catch( ... )
		{
			return false;
		}
		if( !SwitchToWardObject( ward_socket_stream ) )
		{
			delete ward_socket_stream;
			return false;
		}
		return true;
	}
	void Destroy()
	{
		SwitchToSubObject( NULL );
	}
};

#endif
